#include "PDTCodeGraph.h"
#include "PDTNode.h"
#include <saig.h>
#include <abc.h>
#include <abc_global.h>
#include "main.h"
#include "cmd.h"
#include <algorithm>
#include "PDTExpNumBinDerived.h"
#include "PDTNumVar.h"
#include <cmath>
#include <PDTAssert.h>
#include <PDTIfElse.h>
#include <PDTWhile.h>
#include <PDTStatement.h>
#include <PDTListOfStatements.h>
#include <PDTExpBoolBinBoolDerived.h>

ABC_NAMESPACE_IMPL_START

#include "PDTPCTraversal.h"

PDTPCTraversal::PDTPCTraversal(PDTCodeGraph & CG, PDTNumAbcObj * aPC, Id2Pc_t & aId2pc):
    cg(CG), pc(aPC), pc_count(0), id2pc(aId2pc) {
    }

unsigned int  
PDTPCTraversal::start() {
    vector < PDTNode * > & nodeVec = cg.nodes;
    cout << "Starting pc traverse" << endl;
    for (unsigned int i = 0; i < nodeVec.size(); i++) {
        PDTNode * node = nodeVec[i];
        if(node->hasParents())
            continue;

        if (node->traverseEx(*this) == Stop) {
            cerr << "Received Signal Stop from node with nodeId = " 
                << node->getId() << "! Exiting now!" << endl;
            return 0;
        }
    }
    cout << "PC Traverse Done" << endl;
    return pc_count;
}

// PDTAssign
Traversal::RetType
PDTPCTraversal::pre(PDTAssign & d) {
    unsigned int id = d.getId();
    if (id2pc.find(id) != id2pc.end())
        if (id2pc[id]->post != -1)
            return Skip;

    id2pc[id] = new PDTPCHolder();
    id2pc[id]->pre = pc_count++;
    //cout << "In pre, id = " << id << endl;
    return Continue;
}

Traversal::RetType
PDTPCTraversal::post(PDTAssign & d) {
    unsigned int id = d.getId();
    PDTPCHolder * & hptr = id2pc[id];
    if (hptr == NULL)
        return Stop;

    hptr->post = pc_count;
    //cout << "In post, id = " << id << endl;

    stringstream ss;
    hptr->toString(ss);
    cout << ss.str() << endl;

    return Continue; 
}

// PDTPiAssign
Traversal::RetType 
PDTPCTraversal::pre(PDTPiAssign & d) {
    unsigned int id = d.getId();
    if (id2pc.find(id) != id2pc.end())
        if (id2pc[id]->post != -1)
            return Skip;

    id2pc[id] = new PDTPCHolder();
    id2pc[id]->pre = pc_count++;
    return Continue;
}

Traversal::RetType 
PDTPCTraversal::post(PDTPiAssign & d) {
    unsigned int id = d.getId();
    PDTPCHolder * & hptr = id2pc[id];
    if (hptr == NULL)
        return Stop;

    hptr->post = pc_count;

    stringstream ss;
    hptr->toString(ss);
    cout << ss.str() << endl;

    return Continue;
}

// PDTAssert
Traversal::RetType
PDTPCTraversal::pre(PDTAssert & d) {
    unsigned int id = d.getId();
    if (id2pc.find(id) != id2pc.end())
        if (id2pc[id]->post != -1)
            return Skip;

    id2pc[id] = new PDTPCHolder();
    id2pc[id]->pre = pc_count++;
    //cout << "In pre, id = " << id << endl;
    return Continue;
}

Traversal::RetType
PDTPCTraversal::post(PDTAssert & d) {
    unsigned int id = d.getId();
    PDTPCHolder * & hptr = id2pc[id];
    if (hptr == NULL)
        return Stop;

    hptr->post = pc_count;
    //cout << "In post, id = " << id << endl;

    stringstream ss;
    hptr->toString(ss);
    cout << ss.str() << endl;

    return Continue; 
}

// PDTIfElse
Traversal::RetType
PDTPCTraversal::pre(PDTIfElse & d) {
    unsigned int id = d.getId();
    Id2Pc_t::iterator it = id2pc.find(id);
    if (it != id2pc.end()) {
        if (id2pc[id]->post != -1)
            return Skip;
    }

    if (id2pc.find(id) == id2pc.end())
        id2pc[id] = new PDTIfPCHolder();
    id2pc[id]->pre = pc_count++;
    return Continue;
}

Traversal::RetType
PDTPCTraversal::beforeChild(PDTIfElse & d, unsigned int i) {
    PDTNode & child = d.getChildEx(i);
    unsigned int childId = child.getId();
    if (!child.isStatement())
        return Skip;
    if (id2pc[childId] != NULL)
        if(id2pc[childId]->post != -1)
            return Skip;

    if (id2pc[childId] == NULL) {
        if (child.isIfElse())
            id2pc[childId] = new PDTIfPCHolder();
        else if (child.isWhile())
            id2pc[childId] = new PDTWhilePCHolder();
        else
            id2pc[childId] = new PDTPCHolder();
    }

    PDTPCHolder * & child_hld = id2pc[childId];
    PDTIfPCHolder * if_hld = (PDTIfPCHolder*)id2pc[d.getId()];
    if(i == 1) { // if result
        child_hld->pre = pc_count;
        if_hld->begin_then = child_hld->pre;
    } else if (i == 2) { // else result
        child_hld->pre = pc_count;
        if_hld->begin_else = child_hld->pre;
    }
    return Continue;
}

Traversal::RetType 
PDTPCTraversal::afterChild(PDTIfElse & d, unsigned int i) {
    PDTNode & child = d.getChildEx(i);
    if (!child.isStatement())
        return Skip;
    unsigned int childId = child.getId();
    if (id2pc[childId] == NULL)
        return Stop;

    PDTPCHolder * & child_hld = id2pc[childId];
    PDTIfPCHolder * if_hld = (PDTIfPCHolder*)id2pc[d.getId()];
    if(i == 1) {
        if_hld->end_then = child_hld->post;
        pc_count++;
    } else if (i == 2) {
        if_hld->end_else = child_hld->post;
        pc_count++;
    }
    return Continue;
}

Traversal::RetType
PDTPCTraversal::post(PDTIfElse & d) {
    unsigned int id = d.getId();
    if (id2pc.find(id) != id2pc.end())
        if (id2pc[id]->post != -1)
            return Skip;

    PDTIfPCHolder * ifptr = (PDTIfPCHolder*)id2pc[id];

    ifptr->post = pc_count;

    stringstream ss;
    ifptr->toString(ss);
    cout << ss.str() << endl;
    return Continue;
}

// PDTWhile
Traversal::RetType
PDTPCTraversal::pre(PDTWhile & d) {
    unsigned int id = d.getId();
    Id2Pc_t::iterator it = id2pc.find(id);
    if (it != id2pc.end()) {
        if ((*it).second->post != -1)
            return Skip;
    }

    if (it == id2pc.end())
        id2pc[id] = new PDTWhilePCHolder();
    id2pc[id]->pre = pc_count++;
    return Continue;
}

Traversal::RetType
PDTPCTraversal::beforeChild(PDTWhile & d, unsigned int i) {
    PDTNode & child = d.getChildEx(i);
    unsigned int childId = child.getId();
    if (!child.isStatement())
        return Skip;
    if (id2pc[childId] != NULL)
        if(id2pc[childId]->post != -1)
            return Skip;

    if(id2pc[childId] == NULL) {
        if (child.isIfElse())
            id2pc[childId] = new PDTIfPCHolder();
        else if (child.isWhile())
            id2pc[childId] = new PDTWhilePCHolder();
        else
            id2pc[childId] = new PDTPCHolder();
    }

    PDTPCHolder * & child_hld = id2pc[childId];
    PDTWhilePCHolder * while_hld = (PDTWhilePCHolder*)id2pc[d.getId()];
    if (i == 1) { // body statement
        child_hld->pre = pc_count;
        while_hld->begin_while = child_hld->pre;
    }
    return Continue;
}

Traversal::RetType
PDTPCTraversal::afterChild(PDTWhile & d, unsigned int i) {
    PDTNode & child = d.getChildEx(i);
    if (!child.isStatement())
        return Skip;
    unsigned int childId = child.getId();
    if (id2pc[childId] == NULL)
        return Stop;

    PDTPCHolder * & child_hld = id2pc[childId];
    PDTWhilePCHolder * while_hld = (PDTWhilePCHolder*)id2pc[d.getId()];
    if (i == 1) {
        while_hld->end_while = child_hld->post;
        pc_count++;
    }
    return Continue;
}

Traversal::RetType
PDTPCTraversal::post(PDTWhile & d) {
    unsigned int id = d.getId();
    if (id2pc.find(id) != id2pc.end())
        if (id2pc[id]->post != -1)
            return Skip;

    PDTWhilePCHolder * whileptr = (PDTWhilePCHolder *)id2pc[id]; 

    whileptr->post = pc_count;

    stringstream ss;
    whileptr->toString(ss);
    cout << ss.str() << endl;
    return Continue;
}

// PDTListOfStatements
Traversal::RetType 
PDTPCTraversal::pre(PDTListOfStatements & d) {
    //cout << "Entering List of Statements" <<  endl;
    unsigned int id = d.getId();
    if (id2pc[id] != NULL)
        if (id2pc[id]->post != -1)
            return Skip;

    id2pc[id] = new PDTPCHolder();
    id2pc[id]->pre = pc_count;
    return Continue;
}

Traversal::RetType 
PDTPCTraversal::post(PDTListOfStatements & d) {
    //cout << "Exiting List of Statements" <<  endl;
    unsigned int id = d.getId();
    PDTPCHolder * & hptr = id2pc[id];
    if (hptr == NULL)
        return Stop;

    hptr->post = pc_count;
    return Continue;
}

Traversal::RetType 
PDTPCTraversal::beforeChild(PDTListOfStatements & d) {
    return Continue;
}

Traversal::RetType 
PDTPCTraversal::afterChild(PDTListOfStatements & d) {
    return Continue;
}

ABC_NAMESPACE_IMPL_END
